#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2023/10/27 10:42
# @Author  : Hobe Frank
# @Software: PyCharm
import numpy as np
from sklearn.cluster import KMeans
import os

class Cal_pos(object):
    def __init__(self):
        super().__init__()

    def get_pos(self,pixel,ground,width,height):
        f = 102.0  # 如果要改变代码以适应焦距未知，应当把下面xy的值也改变，因为单位要统一，之前我都没考虑到，所以是错的
        suc = False
        xy = (
            np.hstack(
                (
                    0.2 * (pixel[:, 0] - width / 2.0).reshape(-1, 1),
                    0.2 * (-pixel[:, 1] + height / 2.0).reshape(-1, 1),
                )
            )

        )
        xyf = np.hstack((xy, np.full([len(xy), 1], -f)))
        if len(ground) > 10:
            try:
                pos, suc, ground, dis, pho_dis = self.MyRansc(0.95, ground, xyf);
                if not suc:
                    print('定位失败，以上为粗略结果为\n')
            except TypeError:
                suc = False

        if suc:
            return pos, ground, dis, pho_dis
        else:
            pos = np.array([-1])
            ground = [-1]
            dis = -1
            pho_dis = -1
            return pos, ground, dis, pho_dis
    def r_mat(self,alpha, w, k):
        Rf = np.mat([np.cos(alpha), 0, -np.sin(alpha), 0, 1, 0, np.sin(alpha), 0, np.cos(alpha)]).reshape(
            3, 3
        )

        Rw = np.mat([1, 0, 0, 0, np.cos(w), -np.sin(w), 0, np.sin(w), np.cos(w)]).reshape(
            3, 3
        )

        Rk = np.mat([np.cos(k), -np.sin(k), 0, np.sin(k), np.cos(k), 0, 0, 0, 1]).reshape(
            3, 3
        )

        R = Rf * Rw * Rk

        return R

    def calA(self,xyz, uv):
        xyz = xyz.reshape(-1, )
        uv = uv.reshape(-1, )
        Zyiba = R[0, 2] * (xyz[0] - Xs) + R[1, 2] * (xyz[1] - Ys) + R[2, 2] * (xyz[2] - Zs)
        A = np.zeros([2, 7], float);
        A[0, 0] = (R[0, 0] * f + R[0, 2] * (uv[0] - x0)) / Zyiba
        A[0, 1] = (R[1, 0] * f + R[1, 2] * (uv[0] - x0)) / Zyiba
        A[0, 2] = (R[2, 0] * f + R[2, 2] * (uv[0] - x0)) / Zyiba
        A[1, 0] = (R[0, 1] * f + R[0, 2] * (uv[1] - y0)) / Zyiba
        A[1, 1] = (R[1, 1] * f + R[1, 2] * (uv[1] - y0)) / Zyiba
        A[1, 2] = (R[2, 1] * f + R[2, 2] * (uv[1] - y0)) / Zyiba
        A[0, 3] = (uv[1] - y0) * np.sin(w) - (
                (uv[0] - x0) / f * ((uv[0] - x0) * np.cos(k) - (uv[1] - y0) * np.sin(k))
                + f * np.cos(k)
        ) * np.cos(w)
        A[0, 4] = -f * np.sin(k) - (uv[0] - x0) / f * (
                (uv[0] - x0) * np.sin(k) + (uv[1] - y0) * np.cos(k)
        )
        A[0, 5] = uv[1] - y0
        A[1, 3] = -(uv[0] - x0) * np.sin(w) - (
                (uv[1] - y0) / f * ((uv[0] - x0) * np.cos(k) - (uv[1] - y0) * np.sin(k))
                - f * np.sin(k)
        ) * np.cos(w)
        A[1, 4] = -f * np.cos(k) - (uv[1] - y0) / f * (
                (uv[0] - x0) * np.sin(k) + (uv[1] - y0) * np.cos(k)
        )
        A[1, 5] = -(uv[0] - x0)
        A[0, 6] = (uv[0] - x0) / f;
        A[1, 6] = (uv[1] - y0) / f;
        return A

    def cal_uv(self,xyz, uv):
        xyz = xyz.reshape(-1, )
        uv = uv.reshape(-1, )
        x = (
                -f
                * (R[0, 0] * (xyz[0] - Xs) + R[1, 0] * (xyz[1] - Ys) + R[2, 0] * (xyz[2] - Zs))
                / (R[0, 2] * (xyz[0] - Xs) + R[1, 2] * (xyz[1] - Ys) + R[2, 2] * (xyz[2] - Zs))
        )
        y = (
                -f
                * (R[0, 1] * (xyz[0] - Xs) + R[1, 1] * (xyz[1] - Ys) + R[2, 1] * (xyz[2] - Zs))
                / (R[0, 2] * (xyz[0] - Xs) + R[1, 2] * (xyz[1] - Ys) + R[2, 2] * (xyz[2] - Zs))
        )
        l = np.zeros([2, 1], float)
        l[0] = uv[0] - x
        l[1] = uv[1] - y
        return l

    def get_kmeans_random(self,kmeans_list):
        index = np.array([], dtype=np.int64).reshape(-1, 1)
        for i in kmeans_list:
            index = np.concatenate([index, i[np.random.randint(0, len(i), (1, 1))]], axis=0)
        return index

    def cal_distance(self,pois1, pois2):  # 按行储存
        if len(pois1) == len(pois2):
            return np.linalg.norm(pois1 - pois2, axis=1)
        if len(pois1) == 1:
            # pois=pois2.copy();
            pois1 = np.array(list(map(lambda x: pois1, pois2))).reshape(-1, 3)  # 把pois2替换为pois1
            return np.linalg.norm(pois1 - pois2, axis=1)
        if len(pois2) == 1:
            pois2 = np.array(list(map(lambda x: pois2, pois1)))
            return np.linalg.norm(pois1 - pois2, axis=1)
        print("维数错误")
        return False

    def VectorAngle(self,camera_world, ground_world, pos):  # 计算向量夹角，坐标按列储存
        v1 = camera_world - ground_world
        #
        v2 = pos - ground_world
        if len(v1) != 3 or len(v2) != 3:
            print("维数错误")
            exit()
        inner = np.sum(v1 * v2, axis=0).reshape(1, -1)
        # 内积
        mod = (np.linalg.norm(v1, axis=0) * np.linalg.norm(v2, axis=0)).reshape(1, -1)
        # 模的积
        return np.arccos(inner / mod)
        # deg=np.degrees(angle);

    def lineTest(self,camera_world, ground_world, pos):  # 共线检测,阳位情况下像片在光心和地面点之间，如果三点共线则夹角接近180°
        angle = self.VectorAngle(camera_world, ground_world, pos)
        # 弧度
        index1 = np.where(angle > np.pi + 1 / 180.0 * np.pi)[1]
        index2 = np.where(np.pi - 1 / 180.0 * np.pi > angle)[1]
        return np.concatenate([index1, index2])

    def get_pois_kmeans(self,pois, clusters=5):
        kmeans = KMeans(n_clusters=clusters)
        # 使用 KMeans 模型对数据进行聚类
        kmeans.fit(pois)
        predictions = kmeans.predict(pois)
        pois_sorted = []
        for i in range(clusters):
            index = np.where(predictions == i)[0]
            pois_sorted.append(index)
        return pois_sorted

    def MyRansc(self,p, MatchGround, xyf):
        global R, a, w, k, Xs, Ys, Zs, f, x0, y0
        # 迭代最大次数，每次得到更好的估计会优化iters的数值
        iters = 100
        # 数据和模型之间可接受的差值
        # 最好模型的参数估计和内点数目
        total_inlier = 0
        best_pos = 0
        best_dx = 0
        pretotal = 1
        best_suc = False;
        # kmeans_list = get_pois_kmeans(xyf[:,0:2], clusters=10)
        orgin_xyz = [float(np.mean(MatchGround[:, 0])), float(np.mean(MatchGround[:, 1])),
                     float(MatchGround[np.argmin(MatchGround[:, 2]), 2]) + 2]

        for i in range(iters):
            a = 11
            w = 5
            k = 13
            f = 102.0  # 如果要改变代码以适应焦距未知，应当把下面xy的值也改变，因为单位要统一，之前我都没考虑到，所以是错的
            x0 = 0.0
            y0 = 0.0
            Xs = orgin_xyz[0]
            Ys = orgin_xyz[1]
            Zs = orgin_xyz[2]
            R = self.r_mat(a, w, k)
            index = np.random.randint(0, len(MatchGround), (10, 1));
            pos, dx, suc = self.forward(MatchGround[index], xyf[index, :2]);
            # cam_dis = cal_distance(pos, np.array(orgin_xyz).reshape(1, 3))
            if (suc):
                xyf = np.hstack((xyf[:, 0:2], np.full([len(xyf[:, 0:2]), 1], -f)))
                camera_world = np.array(np.dot(R, np.transpose(xyf)) + np.transpose(pos))
                index = self.lineTest(camera_world, np.transpose(MatchGround), np.transpose(pos))  # 坐标按列储存
                good_ground = np.delete(MatchGround, index, axis=0)
                good_camera = np.delete(np.transpose(camera_world), index, axis=0)
                good_xyf = np.delete(xyf, index, axis=0)
                total_inlier = len(good_ground)
                if total_inlier > pretotal:
                    # mean = np.mean(good_ground, axis=1)
                    mean_xyz1 = [float(np.mean(good_camera[:, 0])), float(np.mean(good_camera[:, 1])),
                                 float(np.mean(good_camera[:, 2]))]
                    mean_xyz2 = [float(np.mean(good_ground[:, 0])), float(np.mean(good_ground[:, 1])),
                                 float(np.mean(good_ground[:, 2]))]
                    dis1 = self.cal_distance(pos, np.array(mean_xyz1).reshape(1, 3))
                    dis2 = self.cal_distance(pos, np.array(mean_xyz2).reshape(1, 3))
                    if dis1 > 1500 and dis2 > 900:
                        continue
                    iters = 5 * np.log(1 - p) / np.log(1 - (total_inlier / len(MatchGround)) ** 2)
                    pretotal = total_inlier
                    best_pos = np.copy(pos)
                    best_ground = np.copy(good_ground)
                    best_dx = np.copy(dx)
                    best_suc = suc
                    best_xyf = np.copy(good_xyf)
                    best_dis = np.copy(dis2)
                    best_cam = np.copy(good_camera)
                    best_pho_dis = np.copy(dis1)
                    # 判断是否当前模型已经符合超过一半的点
                    if pretotal > len(MatchGround) * 0.1:
                        print(f'ransac成功\n{len(MatchGround)}个点中，有{pretotal}个内部点\n')
                        # o3d_draw_lines(best_cam, best_ground, best_pos)
                        break;
        if best_suc:
            # pos, dx, suc = forward(best_ground, best_xyf)
            return best_pos, best_suc, best_ground, best_dis, best_pho_dis;
        return np.array([-1, -1, -1]).reshape(1, 3), 0, -1, -1, -1

    def forward(self,ground, pixel):
        global R, a, w, k, Xs, Ys, Zs, f, x0, y0

        # A=np.array(list(map(calA,ground,pixel))).reshape(-1,9)
        # A=np.mat(np.array(list(map(calA,ground,pixel))).reshape(-1,9));
        # # hom_ground=np.concatenate([np.transpose(ground),npP.full([1,len(ground)],1.0)],axis=0);
        # # K=np.array([])
        # l=np.mat(np.array(list(map(cal_uv,ground,pixel))).reshape(-1,1));

        # dx=np.mat(np.full([9,1], 10.0))
        # dx = np.mat(np.full([6, 1], 10.0))
        out = 10.0
        num = 0
        while out > 0.00001:
            A = np.mat(np.array(list(map(self.calA, ground, pixel))).reshape(-1, 7));
            # A = np.mat(np.array(list(map(calA, ground, pixel))).reshape(-1, 6))
            l = np.mat(np.array(list(map(self.cal_uv, ground, pixel))).reshape(-1, 1))
            dx = (A.T * A).I * A.T * l
            Xs = Xs + dx[0]
            Ys = Ys + dx[1]
            Zs = Zs + dx[2]
            a = a + dx[3]
            w = w + dx[4]
            k = k + dx[5]
            f = f + dx[6];
            # x0=x0+dx[7];
            # y0=y0+dx[8];
            R = self.r_mat(float(a), float(w), float(k))
            out = abs(float(dx[np.argmax(abs(dx))]))
            # out=float(max(dx[3],dx[4],dx[5]));
            num = num + 1
            if num == 40:
                print("不收敛")
                break
        v = A * dx - l
        m0 = np.sqrt(v.T * v / (2 * len(v) - 6))
        aa = 1
        # print('平均值X:%f,Y:%f,Z:%f',float(np.mean(ground[:,0])),float(np.mean(ground[:,1])),float(ground[np.argmin(ground[:,2]),2]))
        print("迭代次数为：", num)
        print(
            "结果(m)：\nXs:%16.6f\nYs:%16.6f\nZs:%16.6f\n" % (float(Xs), float(Ys), float(Zs))
        )
        print("单位权中误差m=±%f(m)" % float(m0))
        if num == 40 or np.isnan(Xs[0, 0]):
            return np.array(np.concatenate([Xs, Ys, Zs]).reshape(1, -1)), dx, 0  # 失败
        return np.array(np.concatenate([Xs, Ys, Zs]).reshape(1, -1)), dx, 1  # 成功

